// Copyright FYT 邹承甫 2022

#include "post_processor/processor_node.hpp"

#include <tuple>

#include <fmt/color.h>
#include <fmt/core.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/TransformStamped.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <visualization_msgs/MarkerArray.h>
#include <opencv2/opencv.hpp>

#include "myutils/utils.hpp"
#include "rm_interfaces/SerialSendMsg.h"

namespace post_processor {
ProcessorNode::ProcessorNode() : yaw_(0), pitch_(0) {
  // 读一些常规参数
  nh_.param("/debug", debug_, false);
  nh_.param("/post_processor/world_frame_id", target_frame_,
            std::string("imu"));
  nh_.param("/post_processor/camera_frame_id", camera_frame_,
            std::string("camera"));
  std::vector<double> translation;
  nh_.getParam("/cam_t_gun", translation);
  utils::toEigenMatrix(translation_, translation);

  // 构造tracker
  std::vector<double> f_vec;
  std::vector<double> p_vec;
  std::vector<double> q_vec;
  std::vector<double> r_vec;
  std::vector<double> h_vec;
  double max_spinning_angle, max_match_distance;
  int max_queue_length;
  nh_.getParam("/post_processor/kalman/Q", q_vec);
  nh_.getParam("/post_processor/kalman/F", f_vec);
  nh_.getParam("/post_processor/kalman/R", r_vec);
  nh_.getParam("/post_processor/kalman/P", p_vec);
  nh_.getParam("/post_processor/kalman/H", h_vec);
  nh_.param("/post_processor/tracker/max_spinning_angle", max_spinning_angle,
            30.0);
  nh_.param("/post_processor/tracker/max_match_distance", max_match_distance,
            500.0);
  nh_.param("/post_processor/tracker/max_queue_length", max_queue_length, 5);

  tracker_ = std::make_unique<Tracker>(q_vec, p_vec, f_vec, h_vec, r_vec,
                                       max_match_distance, max_spinning_angle,
                                       max_queue_length, debug_);

  // 构造predictor
  double air_fraction;
  nh_.param("/post_processor/air_fraction", air_fraction, 0.01);
  predictor_ = std::make_unique<Predictor>(air_fraction);

  // 构造pnp_solver
  std::vector<float> camera_matrix;
  std::vector<float> dist_coeffs;
  while (!nh_.hasParam("/camera/camera_matrix") ||
         !nh_.hasParam("/camera/dist_coeffs")) {
    fmt::print(
        "can't get /camera/camera_matrix or /camera/dist_coeffs from ros "
        "parameter server!\n");
    continue;
  };
  nh_.getParam("/camera/camera_matrix", camera_matrix);
  nh_.getParam("/camera/dist_coeffs", dist_coeffs);
  pnp_solver_ = std::make_unique<PnPSolver>(camera_matrix, dist_coeffs);

  // 话题订阅
  targets_sub_.subscribe(nh_, "/all_targets", 10);
  serial_sub_ =
      nh_.subscribe("/serial_receive", 1, &ProcessorNode::serialCallback, this);
  pub_ =
      nh_.advertise<rm_interfaces::SerialSendMsg>("/serial_sendmsg", 1, true);

  // tf
  tf_buffer_ = std::make_shared<tf2_ros::Buffer>();
  tf_listenner_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
  tf_filter_ = std::make_shared<tf2_filter>(targets_sub_, *tf_buffer_,
                                            target_frame_, 10, nh_);
  tf_filter_->registerCallback(&ProcessorNode::targetsCallback, this);
  goal_axis_broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>();

  // 构造动态参数服务器
  f_ = boost::bind(&ProcessorNode::configure_cb, this, _1, _2);
  config_server_.setCallback(f_);

  // 初始化可视化marker
  if (debug_) {
    initMarkers();
  }
}  // end constructor
void ProcessorNode::serialCallback(
    const rm_interfaces::SerialReceiveMsgConstPtr &msg) {
  predictor_->set_bullet_speed(msg->bullet_speed);
  yaw_ = msg->yaw;
  pitch_ = msg->pitch;
}

void ProcessorNode::targetsCallback(
    const rm_interfaces::AllTargetsConstPtr &targets) {
  time_now_ = targets->header.stamp;

  // 解算装甲板的世界坐标
  std::vector<Armor> armors_solved = solveArmorPoses(targets);
  // 跟踪器获得攻击目标
  Eigen::VectorXd target_state = tracker_->update(armors_solved);
  if (target_state.norm() == 0) {
    // 没找到要瞄准的目标
    if (debug_) {
      fmt::print(fg(fmt::color::yellow) | fmt::emphasis::italic,
                 "No target is needed to track\n");
      publishEmptyMarkers();
      publishGoalAxisTf(0, 0);
    }
    publishMsg(0, 0, -1);
    return;
  }
  if (debug_) {
    fmt::print(fg(fmt::color::light_blue) | fmt::emphasis::italic |
                   fmt::emphasis::bold,
               "Tracking target:\n  Position:[{:.1f}, {:.1f}, {:.1f}]mm, \n  "
               "Velocity:[{:.1f}, {:.1f}, {:.1f}]mm, \n  Angle:{:.1f}°\n",
               target_state[0], target_state[1], target_state[2],
               target_state[3], target_state[4], target_state[5],
               target_state[6]);
  }
  // 目标预测
  auto [goal_pitch, predict_state] = predictor_->predict(target_state);
  // 求俯仰偏航
  auto goal_yaw = static_cast<float>(atan2(predict_state[1], predict_state[0]) *
                                     180.0f / CV_PI);
  auto distance = static_cast<float>(sqrt(pow(predict_state[0], 2) +
                                          pow(predict_state[1], 2) +
                                          pow(predict_state[2], 2)) /
                                     1000.0f);
  if (!predictor_->ok()) {
    auto xy_distance = static_cast<float>(
        sqrt(pow(target_state[0], 2) + pow(target_state[1], 2)));
    goal_pitch = static_cast<float>((atan2(target_state[2], xy_distance)) *
                                    180.0f / CV_PI);
  }
  // 话题发送
  publishMsg(goal_yaw, goal_pitch, distance);

  // 可视化目标点
  if (debug_) {
    publishMarkers(predict_state, target_state, armors_solved);
    publishGoalAxisTf(goal_yaw, goal_pitch);
    ros::Time process_end_time = ros::Time::now();
    fmt::print(fg(fmt::color::antique_white) | fmt::emphasis::italic,
               "Total latency is {:.2f}ms\n",
               (process_end_time - time_now_).toSec() * 1000);
  }
}

std::vector<Armor> ProcessorNode::solveArmorPoses(
    const rm_interfaces::AllTargetsConstPtr &targets) {
  std::vector<Armor> armors_solved;
  if (!targets->armors.empty()) {
    for (auto armor : targets->armors) {
      // 不打工程
      // if (armor.id == 2) {
      //   continue;
      // }
      // pnp解算出相机坐标
      std::vector<cv::Point2f> img_pts{{armor.tl.x, armor.tl.y},
                                       {armor.bl.x, armor.bl.y},
                                       {armor.br.x, armor.br.y},
                                       {armor.tr.x, armor.tr.y}};
      bool is_big = armor.is_big || (armor.id == 1) || (armor.id == 8);
      Eigen::VectorXd pose = pnp_solver_->solvePnP(img_pts, is_big);
      if (pose.norm() == 0) {
        continue;
      }
      geometry_msgs::PointStamped camera_pos;
      camera_pos.header = targets->header;
      camera_pos.header.frame_id = camera_frame_;
      camera_pos.point.x = pose[0] + translation_[0];
      camera_pos.point.y = pose[1] + translation_[1];
      camera_pos.point.z = pose[2] + translation_[2];
      geometry_msgs::Point world_pos;
      try {
        world_pos = tf_buffer_->transform(camera_pos, target_frame_).point;
      } catch (const tf2::ExtrapolationException &ex) {
        ROS_WARN("Tf transform failed!\n");
        continue;
      }
      // 计算装甲板角度
      double armor_angle = atan2(pose[3], pose[5]) * 180.0 / CV_PI;
      // 压入Armors
      Eigen::Vector3d target_world(world_pos.x, world_pos.y, world_pos.z);
      armors_solved.emplace_back(target_world, time_now_, armor.id,
                                 armor_angle);
    }
  }
  return armors_solved;
}

void ProcessorNode::publishMsg(const float goal_yaw, const float goal_pitch,
                               const float distance) {
  rm_interfaces::SerialSendMsg msg;
  msg.yaw = goal_yaw - yaw_;
  msg.pitch = goal_pitch - pitch_;
  msg.distance = distance;
  if (debug_) {
    fmt::print(
        fg(fmt::color::lawn_green) | fmt::emphasis::italic |
            fmt::emphasis::bold,
        "Send control cmd [yaw:{:.3f}°, pitch:{:.3f}°, distance:{:.3f}m]\n\n",
        msg.yaw, msg.pitch, msg.distance);
  }
  pub_.publish(msg);
}

void ProcessorNode::publishGoalAxisTf(double yaw, double pitch) {
  geometry_msgs::TransformStamped transform_stamped;
  transform_stamped.header.stamp = ros::Time::now();
  transform_stamped.header.frame_id = target_frame_;
  transform_stamped.child_frame_id = "goal";
  transform_stamped.transform.translation.x = 0.0;
  transform_stamped.transform.translation.y = 0.0;
  transform_stamped.transform.translation.z = 0.0;
  tf2::Quaternion q;
  yaw = yaw / 180.0 * CV_PI;
  pitch = pitch / 180.0 * CV_PI;
  q.setRPY(0, -pitch, yaw);
  transform_stamped.transform.rotation.x = q.x();
  transform_stamped.transform.rotation.y = q.y();
  transform_stamped.transform.rotation.z = q.z();
  transform_stamped.transform.rotation.w = q.w();

  goal_axis_broadcaster_->sendTransform(transform_stamped);
}

void ProcessorNode::initMarkers() {
  vis_pub_ =
      nh_.advertise<visualization_msgs::MarkerArray>("visualization_marker", 0);
  position_marker_.ns = "position";
  position_marker_.type = visualization_msgs::Marker::SPHERE_LIST;
  position_marker_.scale.x = position_marker_.scale.y =
      position_marker_.scale.z = 0.15;
  position_marker_.color.a = 1.0;
  position_marker_.color.r = 1.0;
  predict_marker_.ns = "predict";
  predict_marker_.type = visualization_msgs::Marker::SPHERE;
  predict_marker_.scale.x = predict_marker_.scale.y = predict_marker_.scale.z =
      0.15;
  predict_marker_.color.a = 1.0;
  predict_marker_.color.g = 1.0;
  velocity_marker_.type = visualization_msgs::Marker::ARROW;
  velocity_marker_.ns = "velocity";
  velocity_marker_.scale.x = 0.03;
  velocity_marker_.scale.y = 0.05;
  velocity_marker_.color.a = 1.0;
  velocity_marker_.color.b = 1.0;
  velocity_marker_.color.g = 0.2;
  aim_shoot_marker_.type = visualization_msgs::Marker::ARROW;
  aim_shoot_marker_.ns = "aim_shot";
  aim_shoot_marker_.scale.x = 0.03;
  aim_shoot_marker_.scale.y = 0.05;
  aim_shoot_marker_.color.a = 1.0;
  aim_shoot_marker_.color.g = 1.0;
  aim_shoot_marker_.color.r = 1.0;
  text_marker_.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
  text_marker_.scale.z = 0.15;
  text_marker_.color.a = 1.0;
  text_marker_.color.r = 1.0;
  text_marker_.color.g = 1.0;
  text_marker_.color.b = 1.0;
  text_marker_.lifetime = ros::Duration(0.1);
}

void ProcessorNode::publishEmptyMarkers() {
  position_marker_.action = visualization_msgs::Marker::DELETE;
  velocity_marker_.action = visualization_msgs::Marker::DELETE;
  predict_marker_.action = visualization_msgs::Marker::DELETE;
  text_marker_.action = visualization_msgs::Marker::DELETE;
  visualization_msgs::MarkerArray marker_array;
  marker_array.markers.emplace_back(position_marker_);
  marker_array.markers.emplace_back(predict_marker_);
  marker_array.markers.emplace_back(velocity_marker_);
  marker_array.markers.emplace_back(aim_shoot_marker_);
  marker_array.markers.emplace_back(text_marker_);
  vis_pub_.publish(marker_array);
}

void ProcessorNode::publishMarkers(const Eigen::VectorXd &predict_state,
                                   const Eigen::VectorXd &track_state,
                                   const std::vector<Armor> &armors) {
  visualization_msgs::MarkerArray marker_array;
  std_msgs::Header marker_header;
  marker_header.stamp = time_now_;
  marker_header.frame_id = target_frame_;
  position_marker_.header = predict_marker_.header = velocity_marker_.header =
      aim_shoot_marker_.header = text_marker_.header = marker_header;

  position_marker_.action = visualization_msgs::Marker::ADD;
  position_marker_.points.clear();
  text_marker_.action = visualization_msgs::Marker::ADD;
  text_marker_.id = 0;
  for (auto &armor : armors) {
    geometry_msgs::Point pt;
    pt.x = armor.position[0] / 1000.0;
    pt.y = armor.position[1] / 1000.0;
    pt.z = armor.position[2] / 1000.0;
    position_marker_.points.emplace_back(pt);
    ++text_marker_.id;
    text_marker_.pose.position = pt;
    text_marker_.pose.position.y -= 0.1;
    text_marker_.text = "id:" + std::to_string(armor.id) + "    angle:" +
                        std::to_string(static_cast<int>(armor.yaw));
    marker_array.markers.emplace_back(text_marker_);
  }
  geometry_msgs::Point pt;
  pt.x = track_state[0] / 1000.0;
  pt.y = track_state[1] / 1000.0;
  pt.z = track_state[2] / 1000.0;

  predict_marker_.action = visualization_msgs::Marker::ADD;
  predict_marker_.pose.position.x = predict_state[0] / 1000.0;
  predict_marker_.pose.position.y = predict_state[1] / 1000.0;
  predict_marker_.pose.position.z = predict_state[2] / 1000.0;

  velocity_marker_.action = visualization_msgs::Marker::ADD;
  velocity_marker_.points.clear();
  velocity_marker_.points.emplace_back(pt);
  pt.x += predict_state[3] / 2000.0;
  pt.y += predict_state[4] / 2000.0;
  pt.z += predict_state[5] / 2000.0;
  velocity_marker_.points.emplace_back(pt);

  aim_shoot_marker_.header.frame_id = camera_frame_;
  aim_shoot_marker_.action = visualization_msgs::Marker::ADD;
  aim_shoot_marker_.points.clear();

  aim_shoot_marker_.points.emplace_back();
  geometry_msgs::Point aim_end;
  aim_end.z = 10;
  aim_shoot_marker_.points.emplace_back(aim_end);

  marker_array.markers.emplace_back(position_marker_);
  marker_array.markers.emplace_back(predict_marker_);
  marker_array.markers.emplace_back(velocity_marker_);
  marker_array.markers.emplace_back(aim_shoot_marker_);
  vis_pub_.publish(marker_array);
}

void ProcessorNode::configure_cb(const processor::ProcessorConfig &config,
                                 uint32_t level) {
  predictor_->set_air_fraction(config.air_fraction);
  tracker_->set_predict_time(config.manual_time_offset);
}

}  // namespace post_processor